# -*- coding: utf-8 -*-
"""
项目配置文件
统一管理Hugging Face镜像和模型缓存配置

作者: [您的姓名]
日期: 2024
"""

import os
from pathlib import Path
from loguru import logger

# 项目根目录
PROJECT_ROOT = Path(__file__).parent.absolute()

# Hugging Face 镜像配置
class HuggingFaceConfig:
    """
    Hugging Face 配置类
    统一管理镜像源和缓存目录
    """
    
    # 国内镜像源配置
    MIRRORS = {
        'hf-mirror': 'https://hf-mirror.com',  # 推荐：HF Mirror（稳定快速）
        'modelscope': 'https://www.modelscope.cn',  # 阿里云魔搭社区
        'official': 'https://huggingface.co',  # 官方源（需要VPN）
    }
    
    # 默认使用的镜像
    DEFAULT_MIRROR = 'hf-mirror'
    
    # 模型缓存目录
    MODEL_CACHE_DIR = PROJECT_ROOT / 'models_cache'
    
    @classmethod
    def setup_mirror(cls, mirror_name: str = None):
        """
        设置Hugging Face镜像源
        
        Args:
            mirror_name: 镜像名称，可选值：'hf-mirror', 'modelscope', 'official'
                        如果为None，则使用默认镜像
        """
        if mirror_name is None:
            mirror_name = cls.DEFAULT_MIRROR
            
        if mirror_name not in cls.MIRRORS:
            raise ValueError(f"不支持的镜像源: {mirror_name}，支持的镜像源: {list(cls.MIRRORS.keys())}")
        
        mirror_url = cls.MIRRORS[mirror_name]
        
        # 设置环境变量
        if mirror_name == 'hf-mirror':
            os.environ['HF_ENDPOINT'] = mirror_url
        elif mirror_name == 'modelscope':
            # ModelScope 需要特殊配置
            os.environ['MODELSCOPE_CACHE'] = str(cls.MODEL_CACHE_DIR)
        
        # 设置缓存目录
        cls.setup_cache_dirs()
        
        logger.info(f"✅ 已设置Hugging Face镜像源: {mirror_name} ({mirror_url})")
        logger.info(f"📁 模型缓存目录: {cls.MODEL_CACHE_DIR}")
    
    @classmethod
    def setup_cache_dirs(cls):
        """
        设置模型缓存目录
        """
        # 确保缓存目录存在
        cls.MODEL_CACHE_DIR.mkdir(exist_ok=True)
        
        # 设置各种缓存环境变量
        os.environ['SENTENCE_TRANSFORMERS_HOME'] = str(cls.MODEL_CACHE_DIR)
        os.environ['TRANSFORMERS_CACHE'] = str(cls.MODEL_CACHE_DIR)
        os.environ['HF_HOME'] = str(cls.MODEL_CACHE_DIR)
        
    @classmethod
    def get_current_mirror(cls):
        """
        获取当前使用的镜像源
        
        Returns:
            当前镜像源名称和URL
        """
        hf_endpoint = os.environ.get('HF_ENDPOINT')
        modelscope_cache = os.environ.get('MODELSCOPE_CACHE')
        
        if hf_endpoint:
            for name, url in cls.MIRRORS.items():
                if url == hf_endpoint:
                    return name, url
        elif modelscope_cache:
            return 'modelscope', cls.MIRRORS['modelscope']
        
        return 'official', cls.MIRRORS['official']
    
    @classmethod
    def check_mirror_status(cls):
        """
        检查当前镜像配置状态
        """
        current_mirror, current_url = cls.get_current_mirror()
        cache_dir = cls.MODEL_CACHE_DIR
        
        logger.info("=== Hugging Face 镜像配置状态 ===")
        logger.info(f"当前镜像源: {current_mirror}")
        logger.info(f"镜像URL: {current_url}")
        logger.info(f"缓存目录: {cache_dir}")
        logger.info(f"缓存目录存在: {'是' if cache_dir.exists() else '否'}")
        
        # 检查环境变量
        env_vars = {
            'HF_ENDPOINT': os.environ.get('HF_ENDPOINT', '未设置'),
            'SENTENCE_TRANSFORMERS_HOME': os.environ.get('SENTENCE_TRANSFORMERS_HOME', '未设置'),
            'TRANSFORMERS_CACHE': os.environ.get('TRANSFORMERS_CACHE', '未设置'),
            'HF_HOME': os.environ.get('HF_HOME', '未设置'),
        }
        
        logger.info("\n环境变量:")
        for var, value in env_vars.items():
            logger.info(f"  {var}: {value}")
        
        # 检查可用镜像
        logger.info("\n可用镜像源:")
        for name, url in cls.MIRRORS.items():
            status = "✅ 当前使用" if name == current_mirror else "⭕ 可切换"
            logger.info(f"  {name}: {url} {status}")


# 推荐的模型配置
class ModelConfig:
    """
    推荐的预训练模型配置
    """
    
    # 中文模型推荐
    CHINESE_MODELS = {
        'bert-base-chinese': 'Google官方中文BERT',
        'hfl/chinese-bert-wwm-ext': '哈工大中文BERT（推荐）',
        'hfl/chinese-roberta-wwm-ext': '哈工大中文RoBERTa（推荐）',
        'BAAI/bge-base-zh-v1.5': '智源BGE中文模型（推荐）',
    }
    
    # 多语言模型推荐
    MULTILINGUAL_MODELS = {
        'paraphrase-multilingual-MiniLM-L12-v2': '轻量级多语言模型（推荐）',
        'sentence-transformers/paraphrase-multilingual-mpnet-base-v2': '高性能多语言模型',
        'distiluse-base-multilingual-cased': '蒸馏多语言模型',
    }
    
    # 英文模型推荐
    ENGLISH_MODELS = {
        'all-mpnet-base-v2': '高性能英文模型（推荐）',
        'all-MiniLM-L6-v2': '轻量级英文模型',
        'bert-base-uncased': 'Google官方英文BERT',
    }


# 自动初始化配置
def init_config(mirror_name: str = None, verbose: bool = True):
    """
    初始化项目配置
    
    Args:
        mirror_name: 镜像源名称，默认使用 'hf-mirror'
        verbose: 是否显示详细信息
    """
    try:
        HuggingFaceConfig.setup_mirror(mirror_name)
        if verbose:
            logger.info("\n🎉 项目配置初始化完成！")
            logger.info("💡 提示: 现在可以正常下载和使用Hugging Face模型，无需VPN")
    except Exception as e:
        logger.error(f"❌ 配置初始化失败: {e}")
        raise


# 自动初始化配置，使用国内镜像
# 在模块导入时自动配置，无需手动调用
try:
    HuggingFaceConfig.setup_mirror('hf-mirror')
    logger.info("🌐 已自动配置Hugging Face国内镜像")
except Exception as e:
    logger.warning(f"⚠️ 自动配置镜像失败: {e}")
    logger.info("💡 可以手动调用 init_config() 进行配置")


# 如果直接运行此文件，则显示配置状态
if __name__ == "__main__":
    logger.info("=== QA检索项目配置工具 ===")
    logger.info("\n1. 检查当前配置状态:")
    HuggingFaceConfig.check_mirror_status()
    
    logger.info("\n2. 初始化推荐配置:")
    init_config()
    
    logger.info("\n3. 验证配置:")
    HuggingFaceConfig.check_mirror_status()